﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using System.Diagnostics;

namespace TreeTraversal
{
    public class BalancedBinarySearchTree<T> : BinarySearchTree<T> where T : IComparable
    {
        public override string Name
        {
            get { return "Balanced Binary Search Tree"; }
        }

        public override BinaryTreeNode<T> Add(T Value)
        {
            InternalAddResult result = base.InternalAdd(Value, m_Root);
            BinaryTreeNode<T> tempNode = result.Node;

            if (result.NewNodeAdded)
            {
                while (tempNode != null)
                {
                    BinaryTreeNode<T> nextNode = tempNode.Parent;
                    if (NeedsBalancing(tempNode))
                    {
                        Balance(tempNode);
                    }
                    tempNode = nextNode;
                }
            }

            return result.Node;
        }

        bool NeedsBalancing(BinaryTreeNode<T> Node)
        {
            if (Math.Abs(Node.Balance) > 1)
            {
                return true;
            }
            return false;
        }

        private void Balance(BinaryTreeNode<T> Node)
        {
            int DistanceFromNode = 0;
            BinaryTreeNode<T> RotationNode;
            RotationDirection Direction = RotationDirection.Undefined;
            if (Node.Balance == 0)
            {
                return;
            }
            else if (Node.Balance < 0)
            {
                Direction = RotationDirection.Right;
                // Find right-most node in left node's sub-tree
                RotationNode = Node.Left;
                DistanceFromNode = 1;
                while (RotationNode.Right != null && RotationNode.Balance > 0 && DistanceFromNode < 2)
                {
                    Direction = RotationDirection.Left;
                    RotationNode = RotationNode.Right;
                    DistanceFromNode++;
                }
            }
            else
            {
                Direction = RotationDirection.Left;
                // Find left-most node in right node's sub-tree
                RotationNode = Node.Right;
                DistanceFromNode = 1;
                while (RotationNode.Left != null && RotationNode.Balance < 0 && DistanceFromNode < 2)
                {
                    Direction = RotationDirection.Right;
                    RotationNode = RotationNode.Left;
                    DistanceFromNode++;
                }
            }

            BinaryTreeNode<T> finalNode = Node.Parent;
            while (RotationNode.Parent != finalNode)
            {
                PerformRotation(RotationNode, Direction);
                Direction = ParentDirection(RotationNode);
            }

            // Now that this sub tree is balanced,
            // we need to recalculate all ancestor's depths
            while (Node.Parent != null)
            {
                BinaryTreeNode<T> parentNode = Node.Parent;
                if (parentNode.Left == Node)
                {
                    parentNode.LeftDepth = Node.MaxDepth + 1;
                }
                else if (parentNode.Right == Node)
                {
                    parentNode.RightDepth = Node.MaxDepth + 1;
                }
                else
                {
                    Debug.Assert(false);
                }

                Node = parentNode;

            }
        }

        enum RotationDirection
        {
            Undefined,
            Left,
            Right
        }

        private void PerformRotation(BinaryTreeNode<T> Node, RotationDirection Direction)
        {
            if (Node == null || Node.Parent == null)
                return;

            BinaryTreeNode<T> ParentNode = Node.Parent;
            BinaryTreeNode<T> GrandParentNode = Node.Parent.Parent;
            if (Direction == RotationDirection.Left)
            {
                BinaryTreeNode<T> LeftNode = Node.Left;
                Node.Parent = GrandParentNode;

                ParentNode.Right = LeftNode;
                if (LeftNode != null)
                    LeftNode.Parent = ParentNode;
                ParentNode.RightDepth = (LeftNode != null ? LeftNode.MaxDepth + 1 : 0);

                Node.Left = ParentNode;
                Node.LeftDepth = ParentNode.MaxDepth + 1;
            }
            else if (Direction == RotationDirection.Right)
            {
                BinaryTreeNode<T> RightNode = Node.Right;
                Node.Parent = GrandParentNode;

                ParentNode.Left = RightNode;
                if (RightNode != null)
                    RightNode.Parent = ParentNode;
                ParentNode.LeftDepth = (RightNode != null ? RightNode.MaxDepth + 1 : 0);

                Node.Right = ParentNode;
                Node.RightDepth = ParentNode.MaxDepth + 1;
            }

            ParentNode.Parent = Node;

            if (GrandParentNode != null)
            {
                if (GrandParentNode.Left == ParentNode)
                {
                    GrandParentNode.Left = Node;
                    //GrandParentNode.LeftDepth = Node.MaxDepth + 1;
                }
                else
                {
                    Debug.Assert(GrandParentNode.Right == ParentNode);
                    GrandParentNode.Right = Node;
                    //GrandParentNode.RightDepth = Node.MaxDepth + 1;
                }
            }

            if (Node.Parent == null)
            {
                m_Root = Node;
            }
        }

        RotationDirection ToggleRotation(RotationDirection Rotation)
        {
            if (Rotation == RotationDirection.Left)
                return RotationDirection.Right;
            else if (Rotation == RotationDirection.Right)
                return RotationDirection.Left;
            else
                return RotationDirection.Undefined;
        }

        RotationDirection ParentDirection(BinaryTreeNode<T> Node)
        {
            if (Node.Parent != null)
            {
                if (Node.Parent.Left == Node)
                    return RotationDirection.Right;
                else if (Node.Parent.Right == Node)
                    return RotationDirection.Left;
                else
                    Debug.Assert(false);
            }

            return RotationDirection.Undefined;
        }
    }
}
